from .resnet import resnet50, resnet34, resnet18, resnet101, resnet152
from .resnet_cifar import resnet18_cifar, resnet34_cifar, resnet50_cifar, resnet101_cifar, resnet152_cifar
from .densecl_resnet_cifar import densecl_resnet18_cifar
from .mae_vit_cifar import MAE_ViT_Cifar,MAE_ViT_Classifier
from .mae_vit import mae_vit_base_patch16,mae_vit_huge_patch14,mae_vit_large_patch16
from easydict import EasyDict
import torch



def get_backbone(backbone_name, num_classes=128,**kwargs):
    if backbone_name=="densecl_resnet18_cifar":
        return densecl_resnet18_cifar()
    elif "resnet" in  backbone_name:
        # backbone = resnet50(**kwargs)
        return eval(backbone_name)(num_classes=num_classes)
    elif backbone_name=="mae_cifar_classifier":
        return MAE_ViT_Classifier(num_classes=num_classes,**kwargs)
    else:
        raise NotImplementedError

if __name__=="__main__":
    args=EasyDict()
    args.backbone_name="resnet18"
    x=torch.randn(2,3,224,224)
